import torch
import torch.nn as nn
from mmcls.models import BACKBONES
from mmcls.models.backbones import VisionTransformer
from mmcls.models.utils import resize_pos_embed
from typing import List

import torch.nn.functional as F





@BACKBONES.register_module()
class KVPT(VisionTransformer):
    def __init__(self,
                 prompt_length: int = 1,
                 prompt_layers: List[int] = None,
                 prompt_pos: str = 'postpend',
                 prompt_init: str = 'normal',
                 *args,
                 **kwargs):
        super().__init__(*args, **kwargs)
        for param in self.parameters():
            param.requires_grad = False

        self.prompt_layers = [0] if prompt_layers is None else prompt_layers  
        # self.prompt_layers = range(12) if prompt_layers is None else prompt_layers
        # 577: 576 patches + 1 CLS token
        prompt = torch.empty(len(self.prompt_layers), 577, prompt_length)

        nn.init.normal_(prompt, std=0.02)

        # 在每一层的prompt中添加qkv_extra_params   
        # self.qkv_extra_params = nn.Parameter(qkv_extra_params, requires_grad=True)
        self.prompt = nn.Parameter(prompt, requires_grad=True)
        self.prompt_length = prompt_length
        self.prompt_pos = prompt_pos

    def forward(self, x):
        """Following mmcls implementation."""
        B = x.shape[0] # batch size
        x, patch_resolution = self.patch_embed(x)

        # stole cls_tokens impl from Phil Wang, thanks
        cls_tokens = self.cls_token.expand(B, -1, -1)
        x = torch.cat((cls_tokens, x), dim=1) # x: (B,577,768)
        x = x + resize_pos_embed(
            self.pos_embed,
            self.patch_resolution,
            patch_resolution,
            mode=self.interpolate_mode,
            num_extra_tokens=self.num_extra_tokens) # x: (B,577,768)
        x = self.drop_after_pos(x) # x: (B,577,768)

        # self.prompt: (1=layer,577,1)
        # prompt: (1=layer,4,577,1)
        prompt = self.prompt.unsqueeze(1).expand(-1, x.shape[0], -1, -1)

        outs = []
        for i, layer in enumerate(self.layers):
            if i in self.prompt_layers:
                if self.prompt_pos == 'postpend':
                    # prompt: [layer, batch, length, dim]->(1,4,577,1)
                    ## x: (B,577,768) -> x: (B, 577,769)
                    x_769 = torch.cat([x[:, :, :], prompt[i, :, :, :]], dim=2)
    
            # 重写 x = layer(x)
                def kernel_attention(attention_layer, x):
                    B, N, _ = x.shape # x:(B, 577,769)
                    # 计算核函数
                    v_weight = attention_layer.qkv.weight.view(self.embed_dims*3, self.embed_dims)[2 * self.embed_dims:]

                    # 转置 v_weight 的维度，使其形状变为 (768, 2304 - 2 * 768)
                    v_weight_transposed = v_weight.t()
                    X_flattened = x.view(-1, self.embed_dims)
                    # 计算核函数
                    degree = 2  # Set the degree of the polynomial kernel
                    bias = 1  # Set the bias term
                    kernel_result = torch.pow(torch.matmul(X_flattened, v_weight_transposed) + bias, degree)

                    # 将计算结果还原为原始形状
                    kernel_result = kernel_result.view(B, N, -1)
                    return kernel_result
                # 目的是实现Norm x_769[:,:,self.embed_dims:]@
                kernel_attn = kernel_attention(layer.attn,layer.norm1(x_769[:,:,:self.embed_dims]))
                expanded_vector = x_769[:,:,self.embed_dims:].expand(-1, -1, 768)
                x = x + (kernel_attn * expanded_vector+expanded_vector)
                x = layer.ffn(layer.norm2(x), identity=x)
            else:
                x = layer(x)

            if i == len(self.layers) - 1 and self.final_norm:
                x = self.norm1(x)

            if i in self.out_indices:
                outs.append(x[:, 0])

        return tuple(outs)
